package com.code.train.util;

import cn.hutool.core.lang.UUID;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.lang.management.ManagementFactory;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Redis
 *
 * @param
 * @author 南汉卿
 * @return
 * @date 2025/8/7 15:42
 */
@Slf4j
public class RedisDistributedLock {

    private static final String LOCK_PREFIX = "lock:";
    private static final long DEFAULT_EXPIRE = 30; // 秒
    private static final long RENEW_INTERVAL = 10; // 秒

    private final RedisTemplate<String, String> redisTemplate;

    // 每个线程维护本地缓存的重入次数（默认以线程ID为 ownerId）
    private final ThreadLocal<Map<String, Integer>> reentrantCount = ThreadLocal.withInitial(HashMap::new);

    // 支持传入 ownerId 的可重入计数
    private final ConcurrentMap<String, AtomicInteger> ownerReentrantCount = new ConcurrentHashMap<>();

    // 全局续期调度器
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(4);
    private final ConcurrentMap<String, ScheduledFuture<?>> renewTasks = new ConcurrentHashMap<>();

    // 实例唯一标识，建议用更唯一的 JVM 标识+UUID
    private final String uuid = UUID.randomUUID().toString();

    public RedisDistributedLock(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    // 默认使用线程ID作为 ownerId
    public boolean tryLock(String key) {
        return tryLock(key, getThreadId());
    }

    public boolean tryLock(String key, String ownerId) {
        String redisKey = LOCK_PREFIX + key;
        String lockKey = redisKey + ":" + ownerId;

        // 支持调用链级别的可重入
        AtomicInteger counter = ownerReentrantCount.computeIfAbsent(lockKey, k -> new AtomicInteger(0));
        if (counter.get() > 0) {
            counter.incrementAndGet();
            return true;
        }

        Map<String, Integer> localCountMap = reentrantCount.get();
        if (localCountMap.containsKey(redisKey)) {
            localCountMap.put(redisKey, localCountMap.get(redisKey) + 1);
            return true;
        }

        String luaScript = """
            local owner = redis.call('hget', KEYS[1], 'owner')
            if not owner then
                redis.call('hset', KEYS[1], 'owner', ARGV[1])
                redis.call('hset', KEYS[1], 'count', 1)
                redis.call('expire', KEYS[1], ARGV[2])
                return 1
            elseif owner == ARGV[1] then
                redis.call('hincrby', KEYS[1], 'count', 1)
                redis.call('expire', KEYS[1], ARGV[2])
                return 1
            else
                return 0
            end
        """;

        Boolean success = redisTemplate.execute(
                new DefaultRedisScript<>(luaScript, Boolean.class),
                Collections.singletonList(redisKey),
                ownerId,
                String.valueOf(DEFAULT_EXPIRE)
        );

        if (Boolean.TRUE.equals(success)) {
            counter.incrementAndGet();
            localCountMap.put(redisKey, 1);
            startRenewTask(redisKey, ownerId);
            return true;
        }
        return false;
    }

    // 默认线程释放
    public void unlock(String key) {
        unlock(key, getThreadId());
    }

    public void unlock(String key, String ownerId) {
        String redisKey = LOCK_PREFIX + key;
        String lockKey = redisKey + ":" + ownerId;

        AtomicInteger counter = ownerReentrantCount.get(lockKey);
        if (counter == null || counter.get() == 0) {
            throw new IllegalMonitorStateException("当前ownerId未持有锁，无法释放");
        }

        int newCount = counter.decrementAndGet();
        if (newCount > 0) return;

        String luaScript = """
            local owner = redis.call('hget', KEYS[1], 'owner')
            if owner == ARGV[1] then
                local count = tonumber(redis.call('hget', KEYS[1], 'count'))
                if count > 1 then
                    redis.call('hincrby', KEYS[1], 'count', -1)
                    redis.call('expire', KEYS[1], ARGV[2])
                    return count - 1
                else
                    return redis.call('del', KEYS[1])
                end
            else
                return 0
            end
        """;

        redisTemplate.execute(
                new DefaultRedisScript<>(luaScript, Long.class),
                Collections.singletonList(redisKey),
                ownerId,
                String.valueOf(DEFAULT_EXPIRE)
        );

        reentrantCount.get().remove(redisKey);
        ownerReentrantCount.remove(lockKey);
        stopRenewTask(redisKey);
    }

    private void startRenewTask(String redisKey, String ownerId) {
        renewTasks.computeIfAbsent(redisKey, key -> {
            log.debug("开启续期任务: {}", redisKey);
            return scheduler.scheduleAtFixedRate(() -> {
                try {
                    String luaScript = """
                        local owner = redis.call('hget', KEYS[1], 'owner')
                        if owner == ARGV[1] then
                            return redis.call('expire', KEYS[1], ARGV[2])
                        else
                            return 0
                        end
                    """;

                    redisTemplate.execute(
                            new DefaultRedisScript<>(luaScript, Boolean.class),
                            Collections.singletonList(redisKey),
                            ownerId,
                            String.valueOf(DEFAULT_EXPIRE)
                    );
                    log.debug("成功续期锁: {}", redisKey);
                } catch (Exception e) {
                    log.error("续期任务异常: {}", redisKey, e);
                }
            }, RENEW_INTERVAL, RENEW_INTERVAL, TimeUnit.SECONDS);
        });
    }

    private void stopRenewTask(String redisKey) {
        ScheduledFuture<?> future = renewTasks.remove(redisKey);
        if (future != null) {
            boolean cancelled = future.cancel(true);
            if (cancelled) {
                log.debug("取消续期任务: {}", redisKey);
            }
        }
    }

    private String getThreadId() {
        String jvmId = ManagementFactory.getRuntimeMXBean().getName();
        return uuid + ":" + jvmId + ":" + Thread.currentThread().getId();
    }
}

